package com.zycfc.zsf.boot.web.wrapper;

import javax.servlet.http.*;
import java.io.*;
import javax.servlet.*;

public class ResponseWrapper extends HttpServletResponseWrapper
{
    private ByteArrayOutputStream buffer;
    private ServletOutputStream out;
    
    public ResponseWrapper(final HttpServletResponse httpServletResponse) {
        super(httpServletResponse);
        this.buffer = new ByteArrayOutputStream();
        this.out = new WrapperOutputStream(this.buffer);
    }
    
    public ServletOutputStream getOutputStream() {
        return this.out;
    }
    
    public void flushBuffer() throws IOException {
        if (this.out != null) {
            this.out.flush();
        }
    }
    
    public byte[] getContent() throws IOException {
        this.flushBuffer();
        return this.buffer.toByteArray();
    }
    
    static class WrapperOutputStream extends ServletOutputStream
    {
        private ByteArrayOutputStream bos;
        
        public WrapperOutputStream(final ByteArrayOutputStream bos) {
            this.bos = bos;
        }
        
        public void write(final int b) {
            this.bos.write(b);
        }
        
        public boolean isReady() {
            return false;
        }
        
        public void setWriteListener(final WriteListener arg0) {
        }
    }
}
